# %%
import os
import tensorflow as tf
from tensorflow import keras
# from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
# %%
batch_size = 4
# %%
idg = keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255,
    rotation_range=90,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True,
    vertical_flip=True,
    fill_mode='reflect',
    validation_split=0.1,
)
gx = idg.flow_from_directory(
    'g:/tmp/',
    target_size=(128, 128),
    class_mode=None,
    batch_size=batch_size,
    seed=1,
    interpolation='bicubic',
    subset='training',
)
gy = idg.flow_from_directory(
    'g:/tmp/',
    target_size=(256, 256),
    class_mode=None,
    batch_size=batch_size,
    seed=1,
    interpolation='bicubic',
    subset='training',
)

tx = idg.flow_from_directory(
    'g:/tmp/',
    target_size=(128, 128),
    class_mode=None,
    batch_size=batch_size,
    seed=1,
    interpolation='bicubic',
    subset='validation',
)
ty = idg.flow_from_directory(
    'g:/tmp/',
    target_size=(256, 256),
    class_mode=None,
    batch_size=batch_size,
    seed=1,
    interpolation='bicubic',
    subset='validation',
)
gen = zip(gx, gy)
tgen = zip(tx, ty)
# %%
# dx = gx.next()[0]
# dy = gy.next()[0]
# plt.imshow(dx)
# plt.show()
# plt.imshow(dy)
# plt.show()
# %%
mi = keras.Input((128, 128, 3))
mx = keras.layers.Conv2D(64, 5, padding='same')(mi)
mx = keras.layers.ELU()(mx)
mx = keras.layers.Conv2D(64, 3, padding='same')(mx)
mx = keras.layers.ELU()(mx)
mx = keras.layers.Conv2D(32, 3, padding='same')(mx)
mx = keras.layers.ELU()(mx)
my = keras.layers.Conv2D(3 * (2**2), 3, padding='same')(mx)
mx = keras.layers.ELU()(mx)
my = tf.nn.depth_to_space(my, 2)
model = keras.models.Model(inputs=[mi], outputs=[my])
model.compile('adam', 'mse')
model.summary()
# %%
cp = keras.callbacks.ModelCheckpoint(
    'tmp/weights/upscale_weights.e{epoch:03d}.l{loss:.3f}.vl{val_loss:.3f}.h5',
    save_best_only=True,
    save_weights_only=True,
)
es = keras.callbacks.EarlyStopping(patience=30)
model.fit(gen,
          batch_size=batch_size,
          steps_per_epoch=100,
          initial_epoch=0,
          epochs=500,
          validation_data=tgen,
          validation_steps=10,
          callbacks=[cp, es])

# %%
model.load_weights('tmp/weights/upscale_weights.e022.l0.002.vl0.002.h5')
dx = tx.next()[0]
dy = ty.next()[0]
plt.imshow(dx)
plt.show()
res = model.predict(dx.reshape((-1, 128, 128, 3)))
plt.imshow(res[0])
plt.show()
plt.imshow(dy)
plt.show()
plt.imsave('tmp/test.png', tf.clip_by_value(res[0], 0., 1.).numpy())

# %%
